import cv2
import numpy as np
from ultralytics import YOLO
import supervision as sv

model = YOLO('yolov8n.pt')

VIDEO_PATH = 'videos/e.mp4'

# # 获取视频信息
video_info = sv.VideoInfo.from_video_path(VIDEO_PATH)
# 指定多段线区域
polygons = [
    np.array([[400, 400], [1000, 400], [1000, 700], [400, 700]])
]
zones = [sv.PolygonZone(polygon=polygon, frame_resolution_wh=video_info.resolution_wh) for polygon in polygons]

colors = sv.ColorPalette.default()
# 区域可视化，创建 PolygonZoneAnnotator 对象
zone_annotators = [
    sv.PolygonZoneAnnotator(zone=zone, color=colors.by_idx(0), thickness=6, text_thickness=12, text_scale=4)
    for index, zone in enumerate(zones)
]
# 对目标检对象测框，创建 BoxAnnotator 对象
box_annotators = [
    sv.BoxAnnotator(color=colors.by_idx(2), thickness=2, text_thickness=4, text_scale=2)
    for index in range(len(polygons))
]


# 逐帧处理函数frame: np.ndarray 表示frame是np.ndarray类型， -> np.ndarray返回类型也是np.ndarray
def process_frame(frame: np.ndarray) -> np.ndarray:
    # YOLOV8 推理预测
    results = model(frame, imgsz=1280, verbose=False, show=False, device='cpu')[0]
    # 用 supervision 解析预测结果
    detections = sv.Detections.from_yolov8(results)
    # 遍历每个区域对应的所有 Annotator（其实就是没个端的线）
    for zone, zone_annotator, box_annotator in zip(zones, zone_annotators, box_annotators):
        # 判断目标是否在区域内
        mask = zone.trigger(detections=detections)
        print(mask)
        # 筛选出在区域内的目标
        detections_filtered = detections[mask]
        # 画框
        frame = box_annotator.annotate(scene=frame, detections=detections_filtered, skip_label=True)
        # 画区域，并写区域内目标个数
        frame = zone_annotator.annotate(scene=frame)
    return frame


cap = cv2.VideoCapture(VIDEO_PATH)
while cap.isOpened():
    # Read a frame from the video
    ret, frame = cap.read()
    if not ret:
        break  # 如果没有更多的帧可读取，退出循环
    # 调用 process_frame 方法进行处理
    processed_frame = process_frame(frame, 1)  # 请将 "i" 替换为适当的索引或信息
    # 显示处理后的帧
    cv2.imshow('Processed Frame', processed_frame)
    # 检测按键 "q" 是否被按下，如果是则退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放视频捕捉和视频写入对象
cap.release()
# 关闭所有 OpenCV 窗口
cv2.destroyAllWindows()
